#include "WritePEFile.h"
#include <Windows.h>
#include <time.h>

static UInt RoundToNextPowerOf2(UInt n)
{
	n--;
	n |= n >> 1;
	n |= n >> 2;
	n |= n >> 4;
	n |= n >> 8;
	n |= n >> 16;
	n++;

	return n;
}

UInt elevateToPEPageSizeMultiple(UInt n)
{
	if (n < PE_PAGE_SIZE)
		return PE_PAGE_SIZE;

	return RoundToNextPowerOf2(n);
}

void WritePEFile(FILE* file, struct PEFileCreationParams* params)
{
	struct PEFileSection* s;
	UInt8 blankPage[PE_PAGE_SIZE] = {0};
	UInt i, j;

	// The only two fields that matter are e_magic and e_lfanew
	static const IMAGE_DOS_HEADER dosHdr =
	{
		IMAGE_DOS_SIGNATURE, // Magic number
		0, // Bytes on last page of file
		0, // Pages in file
		0, // Relocations
		0, // Size of header in paragraphs
		0, // Minimum extra paragraphs needed
		0, // Maximum extra paragraphs needed
		0, // Initial (relative) SS value
		0, // Initial SP value
		0, // Checksum
		0, // Initial IP value
		0, // Initial (relative) CS value
		0, // File address of relocation table
		0, // Overlay number
		{0, 0, 0, 0}, // Reserved words
		0, // OEM identifier (for e_oeminfo)
		0, // OEM information; e_oemid specific
		{0, 0, 0, 0, 0, 0, 0, 0, 0, 0}, // Reserved words
		sizeof(IMAGE_DOS_HEADER) // File address of new exe header
	};

	static const UInt32 peSig = IMAGE_NT_SIGNATURE;

	const IMAGE_FILE_HEADER imgFileHdr = 
	{
		IMAGE_FILE_MACHINE_I386, // Machine
		params->numSections, // NumberOfSections
		_time32(NULL), // TimeDateStamp
		0, // PointerToSymbolTable
		0, // NumberOfSymbols
		sizeof(IMAGE_OPTIONAL_HEADER32), // SizeOfOptionalHeader
		IMAGE_FILE_EXECUTABLE_IMAGE | IMAGE_FILE_RELOCS_STRIPPED | IMAGE_FILE_32BIT_MACHINE // Characteristics
	};
	
	IMAGE_OPTIONAL_HEADER32 imgOptHdr =
	{
		0x010B, // Magic
		8, // MajorLinkerVersion
		0, // MinorLinkerVersion
		0 /* TO BE FILLED IN */, // SizeOfCode
		0, // SizeOfInitializedData
		0, // SizeOfUninitializedData
		0 /* TO BE FILLED IN */, // AddressOfEntryPoint
		0 /* TO BE FILLED IN */, // BaseOfCode
		0, // BaseOfData
		0x400000, // ImageBase
		PE_PAGE_SIZE, // SectionAlignment
		PE_PAGE_SIZE, // FileAlignment
		5, // MajorOperatingSystemVersion
		0, // MinorOperatingSystemVersion
		0, // MajorImageVersion
		0, // MinorImageVersion
		4, // MajorSubsystemVersion
		0, // MinorSubsystemVersion
		0, // Win32VersionValue
		0 /* TO BE FILLED IN */, // SizeOfImage
		sizeof(IMAGE_DOS_HEADER) + 4 + sizeof(IMAGE_FILE_HEADER) +
		sizeof(IMAGE_OPTIONAL_HEADER32) + sizeof(IMAGE_SECTION_HEADER), // SizeOfHeaders
		0, // CheckSum
		IMAGE_SUBSYSTEM_WINDOWS_CUI, // Subsystem
		33088, // DllCharacteristics
		0x100000, // SizeOfStackReserve
		0x1000, // SizeOfStackCommit
		0x100000, // SizeOfHeapReserve
		0x1000, // SizeOfHeapCommit
		0, // LoaderFlags
		16, // NumberOfRvaAndSizes
		0 // DataDirectory
	};
	
	IMAGE_SECTION_HEADER sectHdr = 
	{
		{'b', 'l', 'a', 'n', 'k', '\0'}, // Name
		0 /* TO BE FILLED IN */, // VirtualSize
		0 /* TO BE FILLED IN */, // VirtualAddress
		0 /* TO BE FILLED IN */, // SizeOfRawData
		0 /* TO BE FILLED IN */, // PointerToRawData
		0, // PointerToRelocations
		0, // PointerToLinenumbers
		0, // NumberOfRelocations
		0, // NumberOfLinenumbers
		0 /* TO BE FILLED IN */ // Characteristics
	};


	fwrite(&dosHdr, sizeof(IMAGE_DOS_HEADER), 1, file);
	fwrite(&peSig, sizeof(peSig), 1, file);
	fwrite(&imgFileHdr, sizeof(IMAGE_FILE_HEADER), 1, file);

	// Update imgOptHdr
	for (i = 0, j = PE_PAGE_SIZE, s = *params->sections;
		i < params->numSections;
		++i, j += s->numPEPages*PE_PAGE_SIZE, s = params->sections[i])
	{
		if (s->attr & PFSAF_EXECUTE)
		{
			imgOptHdr.SizeOfCode = s->numPEPages * PE_PAGE_SIZE;
			imgOptHdr.AddressOfEntryPoint = j;
			imgOptHdr.BaseOfCode = imgOptHdr.AddressOfEntryPoint;
		}
	}

	imgOptHdr.SizeOfImage = j;

	fwrite(&imgOptHdr, sizeof(IMAGE_OPTIONAL_HEADER32), 1, file);

	// Write section headers
	for (i = 0, j = PE_PAGE_SIZE, s = *params->sections;
		i < params->numSections;
		++i, j += s->numPEPages*PE_PAGE_SIZE, s = params->sections[i])
	{
		sectHdr.Misc.VirtualSize = s->numPEPages * PE_PAGE_SIZE;
		sectHdr.VirtualAddress = j;
		sectHdr.SizeOfRawData = sectHdr.Misc.VirtualSize;
		sectHdr.PointerToRawData = sectHdr.VirtualAddress;
		sectHdr.Characteristics = s->attr;

		fwrite(&sectHdr, sizeof(IMAGE_SECTION_HEADER), 1, file);
	}

	// Align section data by PE_PAGE_SIZE
	i = PE_PAGE_SIZE - ftell(file);
	fwrite(blankPage, i, 1, file);

	// Write section data
	for (i = 0, s = *params->sections; i < params->numSections; ++i, s = params->sections[i])
	{
		fwrite(s->data, s->numPEPages * PE_PAGE_SIZE, 1, file);
	}
}